from transformers import AutoModel
import torch

# 加载 safetensors 格式的模型
model = AutoModel.from_pretrained(r"/root/autodl-tmp/DeepSpeedExamples/applications/fcrlhf/DeepSpeed-Chat/training/step2_reward_model_finetuning/imdb_rm/checkpoint-36000")

# 重新保存为 bin 格式
torch.save(model.state_dict(), "/root/autodl-tmp/DeepSpeedExamples/applications/fcrlhf/DeepSpeed-Chat/training/step2_reward_model_finetuning/imdb_rm/checkpoint-36000/pytorch_model.bin")
